﻿#include "humanseg_c_wrap.h"
#include "humanseg_base.h"
#include "mbv3_humanseg_tflite.h"
#include "pp_humanseg_ort.h"
#include "rvm_humanseg_ort.h"

typedef struct {	
	HUMANSEG_TYPE 	seg_type;
	const char *    model_name;
	const char *	model_path;
	int 			model_width;
	int 			model_height;
}HUMANSEG_MODEL;


const static HUMANSEG_MODEL g_humanseg_models[HUMANSEG_TYPE_MAX] = 
{
	{HUMANSEG_TYPE_PPHUMANSEG_398x224_ORT, "pphumanseg-398x224-ort", "../models/pp_humanseg_398x224_tune1.onnx", 398, 224},
	{HUMANSEG_TYPE_PPHUMANSEG_192x192_ORT, "pphumanseg-192x192-ort", "../models/fcn_hrnetw18_small_v1_192x192_tune1.onnx", 192, 192},
	{HUMANSEG_TYPE_RVM_400x224_ORT, "rvm-400x224-ort", "../models/rvm_mobilenetv3_fp32.onnx", 400, 224},
	{HUMANSEG_TYPE_RVM_512x288_ORT, "rvm-512x288-ort", "../models/rvm_mobilenetv3_fp32.onnx", 512, 288},
	{HUMANSEG_TYPE_RVM_640x360_ORT, "rvm-640x360-ort", "../models/rvm_mobilenetv3_fp32.onnx", 640, 360},	
	{HUMANSEG_TYPE_MBV3_256x144_ORT, "mobilenetv3-256x144-tflite", "../models/mobilenet3.tflite", 256, 144},
};


void humanseg_print_all_models()
{
    int i;

    for(i = 0 ; i < HUMANSEG_TYPE_MAX ; ++i){
        printf("\t%d-%s\n", i, g_humanseg_models[i].model_name);

    }
}

humanseg_ref_t  humanseg_create(int type, int is_video, int enable_debug)
{
    humanseg_ref_t h_seg;

    switch(type){
        case HUMANSEG_TYPE_PPHUMANSEG_398x224_ORT:
        {
            h_seg = (humanseg_ref_t) new PPHumanSegOrt(g_humanseg_models[type].model_path, g_humanseg_models[type].model_width, g_humanseg_models[type].model_height, is_video, enable_debug);
            return h_seg;
        }
        break;
        case HUMANSEG_TYPE_PPHUMANSEG_192x192_ORT:
        {
            h_seg = (humanseg_ref_t) new PPHumanSegOrt(g_humanseg_models[type].model_path, g_humanseg_models[type].model_width, g_humanseg_models[type].model_height, is_video, enable_debug);
            return h_seg;
        }
        break;
        case HUMANSEG_TYPE_RVM_400x224_ORT:
        {
            h_seg = (humanseg_ref_t) new RVMHumanSegOrt(g_humanseg_models[type].model_path, g_humanseg_models[type].model_width, g_humanseg_models[type].model_height, is_video, enable_debug);
            return h_seg;
        }
        break;        
        case HUMANSEG_TYPE_RVM_512x288_ORT:
        {
            h_seg = (humanseg_ref_t) new RVMHumanSegOrt(g_humanseg_models[type].model_path, g_humanseg_models[type].model_width, g_humanseg_models[type].model_height, is_video, enable_debug);
            return h_seg;
        }
        break;
        case HUMANSEG_TYPE_RVM_640x360_ORT:
        {
            h_seg = (humanseg_ref_t) new RVMHumanSegOrt(g_humanseg_models[type].model_path, g_humanseg_models[type].model_width, g_humanseg_models[type].model_height, is_video, enable_debug);
            return h_seg;
        }
        break;

        
        case HUMANSEG_TYPE_MBV3_256x144_ORT:
        {
            h_seg = (humanseg_ref_t) new MBV3HumanSegTFLite(g_humanseg_models[type].model_path, g_humanseg_models[type].model_width, g_humanseg_models[type].model_height, is_video, enable_debug);
            return h_seg;
        }
        break;
        default:
        break;
    };

    return NULL;
}

int humanseg_detect_i420(humanseg_ref_t h_seg, int width, int height, unsigned char **src_plane, int *src_pitch, unsigned char *alpha, int alpha_pitch)
{
    HumanSegBase *p_seg = (HumanSegBase *)h_seg;

    return p_seg->detect_i420(width, height, src_plane, src_pitch, alpha, alpha_pitch);

}

int humanseg_detect_rgb24(humanseg_ref_t h_seg, int width, int height, unsigned char *buf_rgb, int src_pitch, unsigned char *alpha, int alpha_pitch)
{
    HumanSegBase *p_seg = (HumanSegBase *)h_seg;

    return p_seg->detect_rgb24(width, height, buf_rgb, src_pitch, alpha, alpha_pitch);
}

int humanseg_bg_replace_i420(humanseg_ref_t h_seg, int width, int height, unsigned char **fg_plane, int *fg_pitch, 
            unsigned char **bg_plane, int *bg_pitch, unsigned char **mix_plane, int *mix_pitch,
            unsigned char *alpha, int alpha_pitch)
{
    HumanSegBase *p_seg = (HumanSegBase *)h_seg;

    return p_seg->bg_replace_i420(width, height, fg_plane, fg_pitch, 
            bg_plane, bg_pitch, mix_plane, mix_pitch,
            alpha, alpha_pitch);
}

int humanseg_bg_replace_rgb24(humanseg_ref_t h_seg, int width, int height, unsigned char *fg_buf, int fg_pitch, 
            unsigned char *bg_buf, int bg_pitch, unsigned char *mix_buf, int mix_pitch,
            unsigned char *alpha, int alpha_pitch)
{
    HumanSegBase *p_seg = (HumanSegBase *)h_seg;

    return p_seg->bg_replace_rgb24(width, height, fg_buf, fg_pitch, 
            bg_buf, bg_pitch, mix_buf, mix_pitch,
            alpha, alpha_pitch);

}

int humanseg_destroy(humanseg_ref_t h_seg)
{
     HumanSegBase *p_seg = (HumanSegBase *)h_seg;

     delete p_seg;

     return 0;
}